package org.lanyu.springainovel.common.config;

import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.lanyu.springainovel.common.entity.AiModel;
import org.lanyu.springainovel.common.service.AiModelService;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.image.ImageModel;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.OpenAiImageModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.api.OpenAiImageApi;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 模型管理器
 * 用于动态创建和管理不同类型的聊天模型和生图模型
 */
@Slf4j
@Component
@AllArgsConstructor
public class ChatModelManager {

    private final AiModelService aiModelService;

    private final Map<Long, ChatModel> chatModelCache = new ConcurrentHashMap<>();
    private final Map<Long, ImageModel> imageModelCache = new ConcurrentHashMap<>();

    /**
     * 根据模型配置创建聊天模型
     *
     * @param modelConfig 模型配置
     * @return 聊天模型实例
     */
    public ChatModel createChatModel(AiModel modelConfig) {
        if (modelConfig == null) {
            throw new IllegalArgumentException("模型配置不能为空");
        }

        try {
            ChatModel chatModel = switch (modelConfig.getVendor().toLowerCase()) {
                case "ollama" -> createOllamaChatModel(modelConfig);
                default -> createOpenAiChatModel(modelConfig);
            };

            log.info("创建聊天模型成功: {} - {}", modelConfig.getName(), modelConfig.getType());
            return chatModel;
        } catch (Exception e) {
            log.error("创建聊天模型失败: {} - {}", modelConfig.getName(), e.getMessage());
            throw new RuntimeException("创建聊天模型失败", e);
        }
    }

    /**
     * 根据模型配置创建生图模型
     *
     * @param modelConfig 模型配置
     * @return 生图模型实例
     */
    public ImageModel createImageModel(AiModel modelConfig) {
        if (modelConfig == null) {
            throw new IllegalArgumentException("模型配置不能为空");
        }

        try {
            ImageModel imageModel = createOpenAiImageModel(modelConfig);

            log.info("创建生图模型成功: {} - {}", modelConfig.getName(), modelConfig.getType());
            return imageModel;
        } catch (Exception e) {
            log.error("创建生图模型失败: {} - {}", modelConfig.getName(), e.getMessage());
            throw new RuntimeException("创建生图模型失败", e);
        }
    }

    /**
     * 获取或创建聊天模型（带缓存）
     *
     * @param modelConfig 模型配置
     * @return 聊天模型实例
     */
    public ChatModel getOrCreateChatModel(AiModel modelConfig) {
        return chatModelCache.computeIfAbsent(modelConfig.getId(),
                k -> createChatModel(modelConfig));
    }

    /**
     * 获取或创建生图模型（带缓存）
     *
     * @param modelConfig 模型配置
     * @return 生图模型实例
     */
    public ImageModel getOrCreateImageModel(AiModel modelConfig) {
        return imageModelCache.computeIfAbsent(modelConfig.getId(),
                k -> createImageModel(modelConfig));
    }


    /**
     * 根据模型ID查询或创建聊天模型
     * 如果缓存中不存在，则从数据库查询模型配置并创建新的聊天模型
     *
     * @param modelId 模型ID
     * @return 聊天模型实例
     * @throws RuntimeException 如果找不到模型配置或创建模型失败
     */
    public ChatModel getOrCreateChatModelById(Long modelId) {
        // 先从缓存中获取
        ChatModel cachedModel = chatModelCache.get(modelId);
        if (cachedModel != null) {
            return cachedModel;
        }

        // 缓存中不存在，从数据库查询模型配置
        AiModel modelConfig = aiModelService.findById(modelId);
        if (modelConfig == null) {
            throw new RuntimeException("未找到指定模型ID: " + modelId);
        }

        // 创建并缓存新的聊天模型
        return getOrCreateChatModel(modelConfig);
    }

    /**
     * 根据模型ID查询或创建生图模型
     * 如果缓存中不存在，则从数据库查询模型配置并创建新的生图模型
     *
     * @param modelId 模型ID
     * @return 生图模型实例
     * @throws RuntimeException 如果找不到模型配置或创建模型失败
     */
    public ImageModel getOrCreateImageModelById(Long modelId) {
        // 先从缓存中获取
        ImageModel cachedModel = imageModelCache.get(modelId);
        if (cachedModel != null) {
            return cachedModel;
        }

        // 缓存中不存在，从数据库查询模型配置
        AiModel modelConfig = aiModelService.findById(modelId);
        if (modelConfig == null) {
            throw new RuntimeException("未找到指定模型ID: " + modelId);
        }

        // 创建并缓存新的生图模型
        return getOrCreateImageModel(modelConfig);
    }


    /**
     * 清除指定模型的缓�?
     *
     * @param modelId 模型ID
     */
    public void clearCache(Long modelId) {
        chatModelCache.remove(modelId);
        imageModelCache.remove(modelId);
        log.info("清除模型缓存: {}", modelId);
    }

    /**
     * 清除所有缓�?
     */
    public void clearAllCache() {
        chatModelCache.clear();
        imageModelCache.clear();
        log.info("清除所有模型缓存");
    }

    /**
     * 创建Ollama聊天模型
     */
    private ChatModel createOllamaChatModel(AiModel modelConfig) {
        OllamaApi ollamaApi = OllamaApi.builder()
                .baseUrl(modelConfig.getApiUrl())
                .build();

        return OllamaChatModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder()
                        .model(modelConfig.getName())
                        .build())
                .build();
    }

    /**
     * 创建OpenAI聊天模型
     */
    private ChatModel createOpenAiChatModel(AiModel modelConfig) {
        OpenAiApi openAiApi = OpenAiApi.builder()
                .baseUrl(modelConfig.getApiUrl())
                .apiKey(modelConfig.getApiKey())
                .build();

        return OpenAiChatModel.builder()
                .openAiApi(openAiApi)
                .defaultOptions(OpenAiChatOptions.builder()
                        .model(modelConfig.getName())
                        .build())
                .build();
    }

    /**
     * 创建OpenAI生图模型
     */
    private ImageModel createOpenAiImageModel(AiModel modelConfig) {
        OpenAiImageApi openAiImageApi = OpenAiImageApi.builder()
                .baseUrl(modelConfig.getApiUrl())
                .apiKey(modelConfig.getApiKey())
                .build();

        return new OpenAiImageModel(openAiImageApi);
    }
}
